import numpy as np
import random as rd

def load_weights(weights_path):
    '''
    导入训练好的Softmax模型
    :param weights_path: 权重存储位置
    :return: weights:将权重放在矩阵中, m:权重的行数, n:权重的列数
    '''
    f = open(weights_path)
    w = []
    for line in f.readlines():
        w_tmp = []
        lines = line.strip().split("\t")
        for x in lines:
            w_tmp.append(float(x))
        w.append(w_tmp)
    f.close()
    weights = np.mat(w)
    m, n = np.shape(weights)
    return weights, m, n

def load_data(num, m):
    '''
    导入测试数据
    :param num: 生成特是样本的个数
    :param m: 样本的维数
    :return: 生成测试样本
    '''
    testDataSet = np.mat(np.ones((num, m)))
    for i in range(num):
        testDataSet[i, 1] = rd.random() * 6 -3
        testDataSet[i, 2] = rd.random() * 15
    return testDataSet

def predict(test_data, weights):
    '''
    利用训练好的doftmax模型测试数据进行预测
    :param test_data: 测试数据的特征
    :param weights: 模型的权重
    :return: 所属的类别
    '''
    h = test_data * weights
    return h.argmax(axis=1)

def save_result(file_name, result):
    '''
    保存最终的结果
    :param file_name:保存结果的文件名
    :param result: 最终的预测结果
    '''
    f_result = open(file_name, "w")
    m = np.shape(result)[0]
    for i in range(m):
        f_result.write(str(result[i, 0]) + "\n")
    f_result.close()

if __name__ == "__main__":

    print("---- 1.load model")
    w, m, n = load_weights("weights")

    print("---- 2.load data")
    test_data = load_data(4000, m)

    print("---- 3.get Prediction")
    result = predict(test_data, w)

    print("---- 4.save prediction")
    save_result("result", result)